{-# LANGUAGE TypeFamilies               #-}
{-# LANGUAGE GADTs                      #-}
{-# LANGUAGE KindSignatures             #-}
{-# LANGUAGE TypeSynonymInstances       #-}
{-# LANGUAGE FlexibleInstances          #-}
{-# LANGUAGE ScopedTypeVariables        #-}
{-# OPTIONS -Wall #-}

--------------------------------------------------------------------------------
-- |
-- Module      :  Orca.Csound.Monad
-- Copyright   :  (c) Stephen Tetley 2013
-- License     :  BSD3
--
-- Maintainer  :  stephen.tetley@gmail.com
-- Stability   :  unstable
-- Portability :  GHC
--
-- Operational monad.
--
--------------------------------------------------------------------------------

module Orca.Csound.Monad
  (


    Orca
  , execOrca
  , PExpr(..)    
  , IString
  , IBool
  , IInt
  , IFloat

  , TableIx
  , tableIx

  , IRate
  , KRate
  , ARate
  , Sig
  , ISig
  , KSig
  , ASig


  , funCall
  , pfield
  , unaryOp
  , binOp
  , relOp
  , genOpcode0
  , genOpcode1
  , genOpcode1_i
  , genOpcode2
  , genOpcode3
  , genOpcode4

  , genRate             -- Don\'t export class

  , IKA_Signal
  , IK_Signal
  , KA_Signal

  , pureE
  , castI
  , castK
  , castA

  , cond
  , (.==.)
  , (./=.)
  , (.>.)
  , (.<.) 
  , (.>=.)
  , (.<=.) 

  , SMax
  , (.+.)
  , (.-.)
  , (.*.)
  , (./.)
      
  ) where


import Orca.Csound.Syntax
import Orca.Csound.Utils

import Control.Monad.Operational                -- package: operational


import Control.Applicative

data St = St { icount, kcount, acount :: !Int }


data CSI :: * -> * where
  CS_Assign   :: Rhs    -> CSI VarId
  CS_Assign2  :: OpCall -> CSI (VarId,VarId)
  CS_Assign3  :: OpCall -> CSI (VarId,VarId,VarId)
  CS_Assign4  :: OpCall -> CSI (VarId,VarId,VarId,VarId)
  CS_Invoke   :: PrimId -> [Expr] -> CSI ()



type Orca a = Program (CSI) a

execOrca :: Orca a -> [Stmt]
execOrca = toListH . fst3 . (runEval `flip` (St 1 1 1)) . evaluate
  where
    fst3 (a,_,_) = a

-- Note - although Sunroof uses Operational, it runs compile 
-- (aka runOrca) in a standard monad to avoid plumbing in the 
-- eval function.

-- Eval is a Writer-State monad...
--
newtype Eval a = Eval { runEval :: St -> (H Stmt, St, a ) }


instance Functor Eval where
  fmap f ma = Eval $ \s -> let (w,s1,a) = runEval ma s in (w,s1,f a)

instance Applicative Eval where
  pure a    = Eval $ \s -> (emptyH, s, a)
  mf <*> ma = Eval $ \s -> let (w1,s1,f) = runEval mf s
                               (w2,s2,a) = runEval ma s1
                           in (w1 `appendH` w2, s2, f a)

instance Monad Eval where
  return    = pure
  ma >>= k  = Eval $ \s -> let (w1,s1,a) = runEval ma s
                               (w2,s2,b) = runEval (k a) s1
                           in (w1 `appendH` w2, s2, b)


tell :: Stmt -> Eval ()
tell w = Eval $ \s -> (wrapH w, s, ())


evaluate :: Orca a -> Eval a
evaluate = eval . view
  where
    eval :: ProgramView CSI a -> Eval a
    eval (Return x) = return x

    eval (CS_Invoke pid args  :>>= k) = do 
        tell $ InvokeO pid args 
        evaluate $ k ()
 
    eval (CS_Assign rhs :>>= k) = do
        vid <- newVarId $ getRateId rhs
        tell $ Assign vid rhs
        evaluate $ k vid

    eval (CS_Assign2 rhs :>>= k) = do
        v1 <- newVarId $ getRateId rhs
        v2 <- newVarId $ getRateId rhs
        tell $ Assign2 v1 v2 rhs
        evaluate $ k (v1,v2)

    eval (CS_Assign3 rhs :>>= k) = do
        v1 <- newVarId $ getRateId rhs
        v2 <- newVarId $ getRateId rhs
        v3 <- newVarId $ getRateId rhs
        tell $ Assign3 v1 v2 v3 rhs
        evaluate $ k (v1,v2,v3)

    eval (CS_Assign4 rhs :>>= k) = do
        v1 <- newVarId $ getRateId rhs
        v2 <- newVarId $ getRateId rhs
        v3 <- newVarId $ getRateId rhs
        v4 <- newVarId $ getRateId rhs
        tell $ Assign4 v1 v2 v3 v4 rhs
        evaluate $ k (v1,v2,v3,v4)


newVarId :: RateId -> Eval VarId
newVarId rid = Eval $ \s -> case rid of
    A_RATE -> let n = acount s in (emptyH, s{acount = n+1}, mkvar n)
    K_RATE -> let n = kcount s in (emptyH, s{kcount = n+1}, mkvar n)
    I_RATE -> let n = icount s in (emptyH, s{icount = n+1}, mkvar n)
  where
    mkvar i = Local rid i


-- | Phantom type wrapping Expressions.
--
newtype PExpr a = PExpr { getPExpr :: Expr } 


type IString        = PExpr String
type IBool          = PExpr Bool
type IInt           = PExpr Int
type IFloat         = PExpr Float


-- | Separate type for table indices (with are otherwise Int) 
-- disallowing Int operations e.g. Eq, Ord, Num.
--
-- This is analogue to disallowing pointer arithmetic.
-- 
data TableIxTag

type TableIx        = PExpr TableIxTag

tableIx :: Int -> TableIx
tableIx = PExpr . Value I_RATE . Int


data IRate          = IRate deriving (Show)
data KRate          = KRate deriving (Show)
data ARate          = ARate deriving (Show)

type Sig rate       = PExpr rate

type ISig           = Sig IRate
type KSig           = Sig KRate
type ASig           = Sig ARate


instance Num IInt where
  (+) = binOp ADD
  (-) = binOp SUB
  (*) = binOp MUL
  abs _     = error $ "abs - no implementation for IInt"
  signum _  = error $ "signum - no implementation for IInt"
  fromInteger = PExpr . Value I_RATE . Int . fromInteger


instance Num IFloat where
  (+) = binOp ADD
  (-) = binOp SUB
  (*) = binOp MUL
  abs _     = error $ "abs - no implementation for IFloat"
  signum _  = error $ "signum - no implementation for IFloat"
  fromInteger = PExpr . Value I_RATE . Float . fromInteger

instance Fractional IFloat where
  (/)       = binOp DIV
  recip _   = error $ "recip - no implementation for IFLoat"
  fromRational = PExpr . Value I_RATE . Float . fromRational


instance Num ISig where
  (+) = binOp ADD
  (-) = binOp SUB
  (*) = binOp MUL
  abs e     = funCall I_RATE "abs" [e]
  signum _  = error $ "signum - no implementation for ISig"
  fromInteger = PExpr . Value I_RATE . Int . fromInteger


instance Fractional ISig where
  (/)       = binOp DIV
  recip _   = error $ "recip - no implementation for ISig"
  fromRational = PExpr . Value A_RATE . Float . fromRational


instance Num KSig where
  (+) = binOp ADD
  (-) = binOp SUB
  (*) = binOp MUL
  abs e     = funCall K_RATE "abs" [e]
  signum _  = error $ "signum - no implementation for KSig"
  fromInteger = PExpr . Value K_RATE . Int . fromInteger


instance Fractional KSig where
  (/)       = binOp DIV
  recip _   = error $ "recip - no implementation for KSig"
  fromRational = PExpr . Value K_RATE . Float . fromRational



instance Num ASig where
  (+) = binOp ADD
  (-) = binOp SUB
  (*) = binOp MUL
  abs e     = funCall A_RATE "abs" [e]
  signum _  = error $ "signum - no implementation for ASig"
  fromInteger = PExpr . Value A_RATE . Int . fromInteger


instance Fractional ASig where
  (/)       = binOp DIV
  recip _   = error $ "recip - no implementation for ASig"
  fromRational = PExpr . Value A_RATE . Float . fromRational


funCall :: RateId -> String -> [PExpr a] -> PExpr a
funCall rid ss = PExpr . FunCall rid (primId ss) . map getPExpr

pfield :: RateId -> Int -> PExpr a
pfield rid = PExpr . Value rid . Pfield


unaryOp :: UnaryOp -> PExpr a -> PExpr a
unaryOp op a = PExpr $ 
    UnaryE op (getPExpr a)


binOp :: BinOp -> PExpr a -> PExpr a -> PExpr a
binOp op a b = PExpr $ 
    BinaryE op (getPExpr a) (getPExpr b)


relOp :: RelOp -> PExpr a -> PExpr a -> PExpr Bool
relOp op a b = PExpr $ 
    RelE op (getPExpr a) (getPExpr b)




-- No need for RateId 
genOpcode0 :: String -> [Expr] -> Orca ()
genOpcode0 name xs = 
    singleton $ CS_Invoke (primId name) xs


genOpcode1 :: RateId -> String -> [Expr] -> Orca (Sig rate)
genOpcode1 rid name xs = fmap (PExpr . Value rid . Var) $ 
    singleton $ CS_Assign $ OpCallRhs $ OpCall rid (primId name) xs


-- Specailized version for I_RATE values.
--
genOpcode1_i :: String -> [Expr] -> Orca (PExpr a)
genOpcode1_i name xs = fmap (PExpr . Value I_RATE . Var) $ 
    singleton $ CS_Assign $ OpCallRhs $ OpCall I_RATE (primId name) xs


genOpcode2 :: RateId -> String -> [Expr] -> Orca (Sig r1, Sig r1)
genOpcode2 rid name xs = fmap (both (PExpr . Value rid . Var)) $ 
    singleton $ CS_Assign2 $ OpCall rid (primId name) xs
  where
    both f = \(a,b) -> (f a, f b)

genOpcode3 :: RateId -> String -> [Expr] -> Orca (Sig r1, Sig r1, Sig r1)
genOpcode3 rid name xs = fmap (three (PExpr . Value rid . Var)) $ 
    singleton $ CS_Assign3 $ OpCall rid (primId name) xs
  where
    three f = \(a,b,c) -> (f a, f b, f c)

genOpcode4 :: RateId -> String -> [Expr] 
           -> Orca (Sig r1, Sig r1, Sig r1, Sig r1)
genOpcode4 rid name xs = fmap (four (PExpr . Value rid . Var)) $ 
    singleton $ CS_Assign4 $ OpCall rid (primId name) xs
  where
    four f = \(a,b,c,d) -> (f a, f b, f c, f d)




class GenRate rate where
  genRate :: rate -> RateId

instance GenRate IRate where
  genRate _ = I_RATE

instance GenRate KRate where
  genRate _ = K_RATE

instance GenRate ARate where
  genRate _ = A_RATE


class GenRate rate => IKA_Signal rate 

instance IKA_Signal ARate
instance IKA_Signal KRate
instance IKA_Signal IRate

class GenRate rate => IK_Signal rate

instance IK_Signal KRate
instance IK_Signal IRate

class GenRate rate => KA_Signal rate

instance KA_Signal KRate
instance KA_Signal ARate


pureE :: PExpr rate -> Orca (PExpr rate)
pureE e = fmap (PExpr . Value rid . Var) $ 
    singleton $ CS_Assign rhs 
  where
    rhs = PureRhs $ getPExpr e
    rid = getRateId rhs


castA :: Sig a -> Sig ARate
castA = PExpr . Cast A_RATE . getPExpr

castK :: Sig a -> Sig KRate
castK = PExpr . Cast K_RATE . getPExpr

castI :: Sig a -> Sig IRate
castI = PExpr . Cast I_RATE . getPExpr


cond :: IBool -> PExpr a -> PExpr a -> PExpr a
cond test a b = PExpr $ CondE (getPExpr test) (getPExpr a) (getPExpr b)


infix 4 .==.

(.==.) :: Eq a => PExpr a -> PExpr a -> IBool
(.==.) = relOp EQU


infix 4 ./=.

(./=.) :: Eq a => PExpr a -> PExpr a -> IBool
(./=.) = relOp NEQU


infix 4 .>. 

(.>.) :: Eq a => PExpr a -> PExpr a -> IBool
(.>.) = relOp GREATTHAN

infix 4 .<. 

(.<.) :: Eq a => PExpr a -> PExpr a -> IBool
(.<.) = relOp LESSTHAN

infix 4 .>=. 

(.>=.) :: Eq a => PExpr a -> PExpr a -> IBool
(.>=.) = relOp GTEQU

infix 4 .<=. 

(.<=.) :: Eq a => PExpr a -> PExpr a -> IBool
(.<=.) = relOp LTEQU


type family SMax a b :: *

type instance SMax ARate ARate = ARate
type instance SMax ARate KRate = ARate
type instance SMax KRate ARate = ARate
type instance SMax ARate IRate = ARate
type instance SMax IRate ARate = ARate

type instance SMax KRate KRate = KRate
type instance SMax KRate IRate = KRate
type instance SMax IRate KRate = KRate

type instance SMax IRate IRate = IRate


infixl 6 .+.

(.+.) :: c ~ SMax a b => Sig a -> Sig b -> Sig c
(.+.) s1 s2 = PExpr $ BinaryE ADD (getPExpr s1) (getPExpr s2)


infixl 6 .-.

(.-.) :: c ~ SMax a b => Sig a -> Sig b -> Sig c
(.-.) s1 s2 = PExpr $ BinaryE SUB (getPExpr s1) (getPExpr s2)

infixl 7 .*.

(.*.) :: c ~ SMax a b => Sig a -> Sig b -> Sig c
(.*.) s1 s2 = PExpr $ BinaryE MUL (getPExpr s1) (getPExpr s2)

infixl 7 ./.

(./.) :: c ~ SMax a b => Sig a -> Sig b -> Sig c
(./.) s1 s2 = PExpr $ BinaryE DIV (getPExpr s1) (getPExpr s2)


